#!/usr/bin/env python3

import time
import numpy as np
from scipy.spatial.transform import Rotation as R

import rclpy
from rclpy.node import Node
from std_msgs.msg import Float64MultiArray
from rclpy.qos import QoSProfile
from rclpy.callback_groups import ReentrantCallbackGroup


class StewartController(Node):
    def __init__(self):
        super().__init__('main_logic')

        # 状态管理
        self.enabled = False
        self.stopped = False
        self.interpolating = False
        self.pose_queue = []
        self.current_pose = [0.0] * 6  # 初始为零位姿

        # QoS 设置
        qos = QoSProfile(depth=10)

        # 订阅 UDP 节点话题数据（9个float：cmd, subcmd1, subcmd2, x,y,z,r,p,y）
        self.create_subscription(
            Float64MultiArray,
            '/udp_send_data',
            self.command_callback,
            qos)

        # 发布插值后的目标位姿给逆解控制器
        self.pose_pub = self.create_publisher(Float64MultiArray, '/q_target', qos)

        # 发布当前位姿反馈（给 UDP 节点）
        self.feedback_pub = self.create_publisher(Float64MultiArray, '/q_feedback', qos)

        # 定时插值发布器（20Hz）
        self.create_timer(0.05, self.timer_callback, callback_group=ReentrantCallbackGroup())

        # 定时反馈当前位姿（10Hz）
        self.create_timer(0.1, self.feedback_callback)

        self.get_logger().info("Stewart 控制器已启动，等待 UDP 指令...")

    def command_callback(self, msg: Float64MultiArray):
        data = msg.data
        if len(data) < 9:
            self.get_logger().warn("接收到非法命令长度")
            return

        cmd = int(data[0])
        subcmd1 = int(data[1])
        subcmd2 = int(data[2])
        pose = list(data[3:9])

        if cmd == 1:
            self.enabled = True
            self.get_logger().info("✅ 系统已使能")
        elif cmd == 2:
            self.enabled = False
            self.interpolating = False
            self.get_logger().info("� 系统已断电")
        elif cmd == 3:
            if self.enabled:
                self.publish_pose([0.0] * 6)
                self.current_pose = [0.0] * 6
                self.get_logger().info("⏹ 平台回零")
        elif cmd == 4:
            if self.enabled and subcmd1 == 1:
                if subcmd2 == 1:
                    self.get_logger().info("� 进入绝对定位模式")
                    self.start_interpolation_absolute(self.current_pose, pose)
                elif subcmd2 == 2:
                    self.get_logger().info("� 进入相对定位模式")
                    self.start_interpolation_relative(pose)
        elif cmd == 5:
            if self.enabled:
                self.get_logger().info("� 进入保留对准模式")
                self.start_interpolation_absolute(self.current_pose, pose)
        elif cmd == 6:
            self.stopped = True
            self.interpolating = False
            self.get_logger().info("⏸ 插值中止")
        elif cmd == 7:
            self.stopped = False
            self.get_logger().info("▶️ 插值继续")
            if self.pose_queue:
                self.interpolating = True

    def start_interpolation_absolute(self, start_pose, target_pose, steps=100):
        self.pose_queue = []
        for i in range(1, steps + 1):
            interp_pose = [(1 - i / steps) * start_pose[j] + (i / steps) * target_pose[j] for j in range(6)]
            self.pose_queue.append(interp_pose)
        self.interpolating = True
        self.get_logger().info(f"� 开始绝对定位插值，共 {steps} 步")

    def start_interpolation_relative(self, pose_delta, steps=100):
        pos_start = np.array(self.current_pose[:3])
        rpy_start = np.array(self.current_pose[3:])

        delta_pos = np.array(pose_delta[:3])
        delta_rpy = np.array(pose_delta[3:])
        delta_pos_step = delta_pos / steps
        delta_rpy_step = delta_rpy / steps

        self.pose_queue = []

        for i in range(1, steps + 1):
            pos_i = pos_start + i * delta_pos_step
            rpy_i = rpy_start + i * delta_rpy_step
            rot_i = R.from_euler('xyz', rpy_i)
            euler_i = rot_i.as_euler('xyz')  # 可保持统一格式（或原样加法也行）

            pose_i = list(pos_i) + list(euler_i)
            self.pose_queue.append(pose_i)

        self.interpolating = True
        self.get_logger().info(f"� 开始相对定位插值，共 {steps} 步")

    def timer_callback(self):
        if self.interpolating and not self.stopped:
            if self.pose_queue:
                next_pose = self.pose_queue.pop(0)
                self.publish_pose(next_pose)
                self.current_pose = next_pose
            else:
                self.interpolating = False
                self.get_logger().info("✅ 插值完成")

    def publish_pose(self, pose):
        msg = Float64MultiArray()
        msg.data = pose
        self.pose_pub.publish(msg)

    def feedback_callback(self):
        msg = Float64MultiArray()
        msg.data = self.current_pose
        self.feedback_pub.publish(msg)

def main(args=None):
    rclpy.init(args=args)
    node = StewartController()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        node.get_logger().info('� 控制节点终止')
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()